"""
LLM监督微调(SFT)数据集生成脚本
基于检索器数据集生成用于LLM监督微调的问答对数据
使用大语言模型API生成高质量的法律问答数据
"""

from openai import OpenAI
import pandas as pd
from tqdm import tqdm

# 从配置文件导入API密钥和基础URL
from configs.config import API_KEY, BASE_URL, DEFAULT_MODEL

# ==================== 初始化OpenAI客户端 ====================
client = OpenAI(
    api_key=API_KEY,
    base_url=BASE_URL,
)

def generate_with_qwen(prompt, model=None):
    """
    使用Qwen模型生成文本回答

    Args:
        prompt (str): 输入提示词
        model (str, optional): 模型名称，默认使用配置文件中的DEFAULT_MODEL

    Returns:
        str: 模型生成的回答文本

    Note:
        关闭thinking模式以获得更直接的回答
    """
    if model is None:
        model = DEFAULT_MODEL

    messages = [{"role": "user", "content": prompt}]
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        extra_body={"enable_thinking": False}  # 关闭思考模式，直接输出答案
    )

    return completion.choices[0].message.content

def getPrompt(query, doc):
    """
    构建用于生成SFT训练数据的提示词

    Args:
        query (str): 用户查询问题
        doc (str): 相关法律文档内容

    Returns:
        str: 完整的提示词，用于指导模型生成高质量回答

    Note:
        提示词设计重点：
        1. 确保回答严格基于提供的文档
        2. 生成自然流畅的回答风格
        3. 避免引入文档外的信息
    """
    prompt = f"""
你的任务是根据用户的提问和相关文档回答问题，该回答应能准确对应到文档内容，并避免引入不相关内容。请生成符合人类自然搜索习惯的回答。

# 文档说明
- 为法律条文/政府规章类文本，具有以下特点：
  - 包含文件名、章节号、编号、条号等定位信息
  - 表述严谨，多使用"应当""不得""处罚"等术语
  - 可能涉及数据标准（如罚款金额、水位高度等）

# 文档
首先，请仔细阅读以下文档：
<文档内容>
{doc}
</文档内容>

# 用户提问
<用户提问>
{query}
</用户提问>

# 生成要求
在生成回答时，请遵循以下规则：
1. 必须确保回答能准确覆盖文档的核心内容。
2. 输出中只包含回答，不需要前缀、解释或格式说明。
3. 必须模拟真实用户的自然语言习惯，使回答自然流畅。
4. 如果有多个自然表达方式可以作为回答，请选择最常见或最自然的一种。
5. 生成的回答必须要紧密联系文档内容。
6. 禁止引入与文档不相关的背景信息或内容。
"""
    return prompt

# ==================== 主要数据生成逻辑 ====================
print("🚀 开始生成SFT训练数据...")

# 存储生成的数据
data = []

# 读取清理后的检索器数据集
print("📖 读取检索器数据集...")
df = pd.read_csv("data/RetrieverDataset_cleaned.csv")
print(f"📊 数据集大小: {df.shape[0]} 条记录")

# ==================== 批量生成问答对 ====================
for i in tqdm(range(df.shape[0]), desc="生成问答对"):
    row = df.iloc[i]
    query = row["query"]          # 用户查询
    doc = row["positive_doc"]     # 相关法律文档

    # 构建提示词
    prompt = getPrompt(query, doc)

    try:
        # 使用LLM生成回答
        result = generate_with_qwen(prompt)

        # 可选：对结果进行清理（当前注释掉）
        # result = clean_text(result)

        # 保存数据：[查询, 文档, 回答]
        data_record = [query, doc, result]
        data.append(data_record)

        # 每100条记录显示一次进度
        if (i + 1) % 100 == 0:
            print(f"  ✓ 已生成 {i + 1}/{df.shape[0]} 条数据")

    except Exception as e:
        print(f"❌ 第{i}条数据生成失败:")
        print(f"   错误: {e}")
        if 'result' in locals():
            print(f"   结果: {result}")
        print()
        continue

# ==================== 保存生成的数据集 ====================
print("💾 保存SFT数据集...")
columns = ["query", "doc", "answer"]
sft_dataset = pd.DataFrame(data, columns=columns)

# 保存为CSV文件
output_file = "data/LLMDataset_SFT.csv"
sft_dataset.to_csv(output_file, index=False, encoding="utf-8-sig")

print(f"🎉 SFT数据集生成完成！")
print(f"📁 文件保存位置: {output_file}")
print(f"📊 最终数据量: {len(data)} 条记录")